#include "eigen/fashion_mnist_load.h"
#include <fstream>
#include <vector>
#include "swap_endian.h"
#include <iostream>
#include "eigen/normalize.h"
#include "eigen/one_hot.h"

namespace  ldl_eigen
{
FashionMnistLoad::FashionMnistLoad(const std::string& resource_dir) : resource_dir_(resource_dir)
{

}

int64_t FashionMnistLoad::load()
{
    try {
        train_images_ = load_images(resource_dir_ + "train-images-idx3-ubyte");
        train_labels_ = load_labels(resource_dir_ + "train-labels-idx1-ubyte");
        test_images_ = load_images(resource_dir_ + "t10k-images-idx3-ubyte");
        test_labels_ = load_labels(resource_dir_ + "t10k-labels-idx1-ubyte");

        normalized_train_images_ = Normalize::normalize(train_images_);
        one_hot_train_labels_ = OneHot::encode(train_labels_).cast<float>();
        normalized_test_images_ = Normalize::normalize(test_images_);
        one_hot_test_labels_ = OneHot::encode(test_labels_).cast<float>();
    } catch (const std::exception& e) {
        std::cerr << "错误: " << e.what() << "\n";
        return 1;
    }

    return 0;
}

const Eigen::MatrixXi FashionMnistLoad::get_train_images()
{
    return train_images_;
}
const Eigen::VectorXi FashionMnistLoad::get_train_labels()
{
    return train_labels_;
}
const Eigen::MatrixXi FashionMnistLoad::get_test_images()
{
    return test_images_;
}
const Eigen::VectorXi FashionMnistLoad::get_test_labels()
{
    return test_labels_;
}

const Eigen::MatrixXf FashionMnistLoad::get_one_hot_train_labels()
{
    return one_hot_train_labels_;
}
const Eigen::MatrixXf FashionMnistLoad::get_one_hot_test_labels()
{
    return one_hot_test_labels_;
}

const Eigen::MatrixXf FashionMnistLoad::get_normalized_train_images()
{
    return normalized_train_images_;
}
const Eigen::MatrixXf FashionMnistLoad::get_normalized_test_images()
{
    return normalized_test_images_;
}

Eigen::MatrixXi FashionMnistLoad::load_images(const std::string& path, int max_items)
{
    std::ifstream file(path, std::ios::binary);
    if (!file) {
        throw std::runtime_error("无法打开文件: " + path);
    }

    uint32_t magic, num_images, rows, cols;
    file.read(reinterpret_cast<char*>(&magic), 4);
    file.read(reinterpret_cast<char*>(&num_images), 4);
    file.read(reinterpret_cast<char*>(&rows), 4);
    file.read(reinterpret_cast<char*>(&cols), 4);

    // 转换为小端序
    magic = SwapEndian::swap(magic);
    num_images = SwapEndian::swap(num_images);
    rows = SwapEndian::swap(rows);
    cols = SwapEndian::swap(cols);

    if (magic != 0x803) {
        throw std::runtime_error("无效的图像文件格式");
    }

    if (max_items > 0 && num_images > static_cast<uint32_t>(max_items)) {
        num_images = max_items;
    }

    Eigen::MatrixXi images(num_images, rows * cols);
    std::vector<uint8_t> read_data(rows * cols);
    for (uint32_t i = 0; i < num_images; ++i) {
        file.read((char *)read_data.data(), rows * cols);
        for(uint32_t j = 0; j < rows * cols; j++) {
            images(i,j) = read_data[j];
        }
    }

    return images;
}

// 加载标签数据
Eigen::VectorXi FashionMnistLoad::load_labels(const std::string& path, int max_items) {
    std::ifstream file(path, std::ios::binary);
    if (!file) {
        throw std::runtime_error("无法打开文件: " + path);
    }

    uint32_t magic, num_labels;
    file.read(reinterpret_cast<char*>(&magic), 4);
    file.read(reinterpret_cast<char*>(&num_labels), 4);

    // 转换为小端序
    magic = SwapEndian::swap(magic);
    num_labels = SwapEndian::swap(num_labels);

    if (magic != 0x801) {
        throw std::runtime_error("无效的标签文件格式");
    }

    if (max_items > 0 && num_labels > static_cast<uint32_t>(max_items)) {
        num_labels = max_items;
    }

    std::vector<char> read_data(num_labels);
    file.read(read_data.data(), num_labels);
    Eigen::VectorXi labels(num_labels);
    for(uint64_t j = 0; j < num_labels;j++)
    {
        labels(j) = read_data[j];
    }

    return labels;
}

}